import os

import numpy as np
import torchvision.datasets as datasets

from PIL import Image

from utils import process_image


class FaceDataset(datasets.ImageFolder):
    def __init__(self, path, pairs_path, image_size, transform=None):
        super(FaceDataset, self).__init__(path, transform)
        self.image_size = image_size
        self.pairs_path = pairs_path
        self.validation_images = self.get_paths(path)

    @staticmethod
    def read_pairs(pairs_file_path):
        """
        读取 pairs.txt 文件
        :param pairs_file_path: pairs.txt 文件路径
        :return: pair list
        """
        pairs = []
        with open(pairs_file_path, 'r') as f:
            for line in f.readlines()[1:]:
                pair = line.strip().split()
                pairs.append(pair)
        return pairs

    def get_paths(self, dir_path, file_ext="jpg"):
        pairs = self.read_pairs(self.pairs_path)
        skipped_pairs = 0
        path_list, is_same_list = [], []

        for i in range(len(pairs)):
            pair = pairs[i]
            path0, path1, is_same = None, None, None

            if len(pair) == 3:  # 属于同一人的不同图片
                path0 = os.path.join(dir_path, pair[0], f"{int(pair[1]):03d}.{file_ext}")
                path1 = os.path.join(dir_path, pair[0], f"{int(pair[2]):03d}.{file_ext}")
                is_same = True
            elif len(pair) == 4:  # 属于不同人的不同图片
                path0 = os.path.join(dir_path, pair[0], f"{int(pair[1]):03d}.{file_ext}")
                path1 = os.path.join(dir_path, pair[2], f"{int(pair[3]):03d}.{file_ext}")
                is_same = False

            if os.path.exists(path0) and os.path.exists(path1):  # pair 对应的两张图片存在
                path_list.append((path0, path1, is_same))
                is_same_list.append(is_same)
            else:  # 否则跳过并计数
                skipped_pairs += 1
        if skipped_pairs > 0:
            print(f"Skipped {skipped_pairs} image pairs.")
        return path_list

    def __getitem__(self, index):
        path_1, path_2, is_same = self.validation_images[index]
        img1, img2 = Image.open(path_1), Image.open(path_2)
        img1 = process_image(img1, [self.image_size[1], self.image_size[0]])
        img2 = process_image(img2, [self.image_size[1], self.image_size[0]])

        img1, img2 = np.array(img1) / 255, np.array(img2) / 255
        img1 = np.transpose(img1, [2, 0, 1])
        img2 = np.transpose(img2, [2, 0, 1])

        return img1, img2, is_same

    def __len__(self):
        return len(self.validation_images)
